import os
import subprocess
train={}
test={}
path='..'
f=open(path+'/trainLabels.csv')
fo=open('trainLabels.csv','w')
fo.write(f.readline())
f1=open(path+'/sampleSubmission.csv')
fo1=open('sampleSubmission.csv','w')
fo1.write(f1.readline())

for line in f:
    xx=line.split(',')
    label=int(xx[-1])
    if label not in train:
        size=os.path.getsize(path+'/train/'+xx[0][1:-1]+'.asm')
        #if size>1000000:
        #    continue
        train[label]=1
        idx=xx[0][1:-1]
        cmd='cp '+path+'/train/'+idx+'.* train/'
        subprocess.call(cmd,shell=True)
        fo.write(line)
    elif label not in test:
        test[label]=1
        idx=xx[0][1:-1]
        cmd='cp '+path+'/train/'+idx+'.* test/'
        subprocess.call(cmd,shell=True)
        fo1.write(idx+','+','.join(['0' if i!=label-1 else '1' for i in range(9)])+'\n')
    
f.close()
fo.close()
f1.close()
fo1.close()
